import matplotlib.pyplot as plt
from sklearn.decomposition import PCA
from sklearn.datasets import load_iris

if __name__ == '__main__':
    data = load_iris()
    features, labels = data.data, data.target
    pca = PCA(n_components=2)
    reduced_features = pca.fit_transform(features)

    red_x, red_y = [], []
    blue_x, blue_y = [], []
    green_x, green_y = [], []

    for i in range(len(reduced_features)):
        if labels[i] == 0:
            red_x.append(reduced_features[i][0])
            red_y.append(reduced_features[i][1])
        elif labels[i] == 1:
            blue_x.append(reduced_features[i][0])
            blue_y.append(reduced_features[i][1])
        else:
            green_x.append(reduced_features[i][0])
            green_y.append(reduced_features[i][1])

    plt.scatter(red_x, red_y, c='red', marker='x', label='setosa')
    plt.scatter(blue_x, blue_y, c='blue', marker='D', label='versicolor')
    plt.scatter(green_x, green_y, c='green', marker='.', label='virginica')
    plt.legend()
    plt.xlabel('First Principal Component')
    plt.ylabel('Second Principal Component')
    plt.title('PCA for Iris Dataset')
    plt.show()
